package com.sbtr.util;

import java.util.List;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLLimit;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIntegerExpr;
import com.alibaba.druid.sql.ast.expr.SQLNumericLiteralExpr;
import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.ast.statement.SQLUnionQuery;
import com.alibaba.druid.util.JdbcUtils;

/**
 * SQL 工具类
 */
public class SQLReadUtils {
	/**
	 * 重新构建SQL限制查询的结果，只返回对应count条数据
	 */
	public static String limit(String sql, String dbType, int count) {
		List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
		if (stmtList.size() != 1) {
			throw new IllegalArgumentException("sql not support count : " + sql);
		}

		SQLStatement stmt = stmtList.get(0);

		if (!(stmt instanceof SQLSelectStatement)) {
			throw new IllegalArgumentException("sql不是查询语句: " + sql);
		}

		SQLSelectStatement selectStmt = (SQLSelectStatement) stmt;

		return limit(selectStmt.getSelect(), dbType, count);
	}

	public static String limit(SQLSelect select, String dbType, int count) {
		limit(select, dbType, count, true);
		return SQLUtils.toSQLString(select, dbType);
	}

	public static boolean limit(SQLSelect select, String dbType, int count, boolean check) {
		SQLSelectQuery query = select.getQuery();
		if (query instanceof SQLUnionQuery) {
			SQLUnionQuery union = (SQLUnionQuery) query;
			return limitUnion(union, dbType, count, check);
		}

		SQLSelectQueryBlock queryBlock = (SQLSelectQueryBlock) query;
		if (JdbcUtils.MYSQL.equals(dbType)) {
			return limitMySqlQueryBlock(queryBlock, dbType, count, check);
		}

		if (JdbcUtils.HIVE.equals(dbType)) {
			return limitHiveQueryBlock(queryBlock, dbType, count, check);
		}

		throw new UnsupportedOperationException();
	}

	private static boolean limitUnion(SQLUnionQuery queryBlock, String dbType, int count, boolean check) {
		SQLLimit limit = queryBlock.getLimit();
		if (limit != null) {
			if (check && limit.getRowCount() instanceof SQLNumericLiteralExpr) {
				int rowCount = ((SQLNumericLiteralExpr) limit.getRowCount()).getNumber().intValue();
				if (rowCount <= count) {
					return false;
				}
			} else if (limit.getRowCount() instanceof SQLVariantRefExpr) {
				return false;
			}

			limit.setRowCount(new SQLIntegerExpr(count));
		}

		if (limit == null) {
			limit = new SQLLimit();
			limit.setOffset(new SQLIntegerExpr(0));
			limit.setRowCount(new SQLIntegerExpr(count));
			queryBlock.setLimit(limit);
		}
		return true;
	}

	private static boolean limitMySqlQueryBlock(SQLSelectQueryBlock queryBlock, String dbType, int count,
			boolean check) {
		SQLLimit limit = queryBlock.getLimit();
		if (limit != null) {
			if (check && limit.getRowCount() instanceof SQLNumericLiteralExpr) {
				int rowCount = ((SQLNumericLiteralExpr) limit.getRowCount()).getNumber().intValue();
				if (rowCount <= count) {
					return false;
				}
			} else if (limit.getRowCount() instanceof SQLVariantRefExpr) {
				return false;
			}
			limit.setRowCount(new SQLIntegerExpr(count));
		}

		if (limit == null) {
			limit = new SQLLimit();
			limit.setOffset(new SQLIntegerExpr(0));
			limit.setRowCount(new SQLIntegerExpr(count));
			queryBlock.setLimit(limit);
		}
		return true;
	}

	private static boolean limitHiveQueryBlock(SQLSelectQueryBlock queryBlock, String dbType, int count,
			boolean check) {
		SQLLimit limit = queryBlock.getLimit();
		if (limit != null) {
			if (check && limit.getRowCount() instanceof SQLNumericLiteralExpr) {
				int rowCount = ((SQLNumericLiteralExpr) limit.getRowCount()).getNumber().intValue();
				if (rowCount <= count) {
					return false;
				}
			} else if (limit.getRowCount() instanceof SQLVariantRefExpr) {
				return false;
			}
			limit.setRowCount(new SQLIntegerExpr(count));
		}

		if (limit == null) {
			limit = new SQLLimit();
			//limit.setOffset(new SQLIntegerExpr(0));
			limit.setRowCount(new SQLIntegerExpr(count));
			queryBlock.setLimit(limit);
		}
		return true;
	}
}
